from theHarvester.discovery.constants import *
from theHarvester.lib.core import *
import json
import time


class SearchPentestTools:

    def __init__(self, word):
        # Script is largely based off https://pentest-tools.com/public/api_client.py.txt
        self.word = word
        self.key = Core.pentest_tools_key()
        if self.key is None:
            raise MissingKey('PentestTools')
        self.total_results = []
        self.api = f'https://pentest-tools.com/api?key={self.key}'
        self.proxy = False

    async def poll(self, scan_id):
        while True:
            time.sleep(3)
            # Get the status of our scan
            scan_status_data = {
                'op': 'get_scan_status',
                'scan_id': scan_id
            }
            responses = await AsyncFetcher.post_fetch(url=self.api, data=json.dumps(scan_status_data), proxy=self.proxy)
            res_json = json.loads(responses.strip())
            if res_json['op_status'] == 'success':
                if res_json['scan_status'] != 'waiting' and res_json['scan_status'] != 'running':
                    getoutput_data = {
                        'op': 'get_output',
                        'scan_id': scan_id,
                        'output_format': 'json'
                    }
                    responses = await AsyncFetcher.post_fetch(url=self.api,
                                                              data=json.dumps(getoutput_data),
                                                              proxy=self.proxy)

                    res_json = json.loads(responses.strip('\n'))
                    self.total_results = await self.parse_json(res_json)
                    break
            else:
                print(f"Operation get_scan_status failed because: {res_json['error']}. {res_json['details']}")
                break

    @staticmethod
    async def parse_json(json_results):
        status = json_results['op_status']
        if status == 'success':
            scan_tests = json_results['scan_output']['output_json']
            output_data = scan_tests[0]['output_data']
            host_to_ip = [f'{subdomain[0]}:{subdomain[1]}' for subdomain in output_data if len(subdomain) > 0]
            return host_to_ip
        return []

    async def get_hostnames(self) -> list:
        return self.total_results

    async def do_search(self):
        subdomain_payload = {
            'op': 'start_scan',
            'tool_id': 20,
            'tool_params': {
                'target': f'{self.word}',
                'web_details': 'off',
                'do_smart_search': 'off'
            }
        }
        responses = await AsyncFetcher.post_fetch(url=self.api, data=json.dumps(subdomain_payload), proxy=self.proxy)
        res_json = json.loads(responses.strip())
        if res_json['op_status'] == 'success':
            scan_id = res_json['scan_id']
            await self.poll(scan_id)

    async def process(self, proxy=False):
        self.proxy = proxy
        await self.do_search()  # Only need to do it once.
